"""
用户登录之后才能访问接口；根据不同状态码自定义请求的成功和失败；
"""
import os

from locust import HttpUser, task, between

from common.config.conf_sample2 import req_headers, form_data, authorization_headers


# HttpUser 用来给每个用户提供一个client属性(HttpSession的实例)，向目标服务器发送HTTP请求
class QuickstartUser(HttpUser):
    """
    目标url：http://www.datagrand.com/demo/api/nlp/classify
    """
    # 每个任务执行之后的等待时间
    wait_time = between(0.9, 1.2)

    # # 用户启动的时候调用
    @task
    def on_start(self):
        login_data = {
            "username": form_data["username"],
            "password": form_data["password"]
        }
        resp = self.client.post("/demo/api/auth", data=login_data, headers=authorization_headers)
        if resp.ok:
            authorization = resp.content.decode('utf-8')
        else:
            authorization = os.getenv("BUSINESS_AUTHORIZATION")
        # 请求头中携带认证token
        req_headers.setdefault("Authorization", "Bearer " + authorization)

    # 标记用户将要执行的任务
    @task
    def task1(self):
        data = {
            "text": "网易财经3月26日讯。今日沪指高开低走，一路震荡下行，午后跌幅进一步扩大，最终失守3000点"
        }
        self.client.post("/demo/api/nlp/classify", headers=req_headers, data=data)

        # 根据不同状态码自定义请求的成功和失败
        # with self.client.post("/demo/api/nlp/classify", headers=req_headers,
        #                       data=data, catch_response=True) as response:
        #     if response.status_code in (400, 401, 403):
        #         print("响应状态码是", response.status_code)
        #         response.success()


if __name__ == '__main__':
    pass
    # 启动: locust -f locustfiles/sample2.py
